#include "OffLineRenderPCH.h"
#include "OR_ScenePrimitiveAggregateBVH.h"
#include "UT_MathHelper.h"

using namespace Utility;
using namespace Memory;
using namespace OffLineRender;

//-----------------------------------------------------------------
//OR_ScenePrimitiveAggregateBVH::BVHPrimitiveInfo
//-----------------------------------------------------------------
struct OR_ScenePrimitiveAggregateBVH::BVHPrimitiveInfo : public MM_DefaultObject
{
public:
	BVHPrimitiveInfo() { }
	BVHPrimitiveInfo(int pn, const UT_BBOX_AABB& b)
		: m_iPrimitiveNumber(pn), m_kWorldBounds(b)
	{
		m_kCentroid = .5f * b.pMin + .5f * b.pMax;
	}

	int32_t			m_iPrimitiveNumber;
	UT_Point3Float	m_kCentroid;
	UT_BBOX_AABB	m_kWorldBounds;
};
//-----------------------------------------------------------------
//OR_ScenePrimitiveAggregateBVH::BVHBuildNode 
//-----------------------------------------------------------------
struct OR_ScenePrimitiveAggregateBVH::BVHBuildNode : public MM_DefaultObject
{
public:
	// BVHBuildNode Public Methods
	BVHBuildNode()
	{ 
		m_akChildrenNode[0] = NULL;
		m_akChildrenNode[1] = NULL; 
	}

	void InitAsLeaf(uint32_t uiFirstPrimOffset, uint32_t uiPrimitiveNums, const UT_BBOX_AABB& rkWorldBound)
	{
		m_uiFirstPrimOffset = uiFirstPrimOffset;
		m_uiPrimitiveNums	= uiPrimitiveNums;
		m_kWorldBounds		= rkWorldBound;
	}

	void InitAsInterior(uint32_t axis, BVHBuildNode *c0, BVHBuildNode *c1)
	{
		m_akChildrenNode[0] = c0;
		m_akChildrenNode[1] = c1;
		m_kWorldBounds		= UT_MathHelper::Union(c0->m_kWorldBounds, c1->m_kWorldBounds);
		m_uiSplitAxis		= axis;
		m_uiPrimitiveNums	= 0;
	}

	UT_BBOX_AABB	m_kWorldBounds;
	BVHBuildNode*	m_akChildrenNode[2];
	uint32_t		m_uiSplitAxis, m_uiFirstPrimOffset, m_uiPrimitiveNums;
};
//-----------------------------------------------------------------
//OR_ScenePrimitiveAggregateBVH::LinearBVHNode
//-----------------------------------------------------------------
struct OR_ScenePrimitiveAggregateBVH::LinearBVHNode : public MM_DefaultObject
{
public:
	UT_BBOX_AABB m_kWorldBounds;

	union
	{
		uint32_t m_uiPrimitivesOffset;    // leaf
		uint32_t m_uiSecondChildOffset;   // interior
	};

	uint8_t m_ucPrimitives;   // 0 -> interior node
	uint8_t m_ucAxis;         // interior node: xyz
	uint8_t m_ucPad[2];       // ensure 32 byte total size
};
//-----------------------------------------------------------------
// OR_ScenePrimitiveAggregateBVH
//-----------------------------------------------------------------
OR_ScenePrimitiveAggregateBVH::OR_ScenePrimitiveAggregateBVH(const Utility::UT_TVector<OR_ScenePrimitiveBasePtr >& p, 
	uint32_t maxPrims /* = 1 */)
{
	m_uiMaxPrimsInNode = UT_MathHelper::Min(255, maxPrims);

	for (uint32_t i = 0; i < p.GetSize(); ++i)
		p.GetAt(i)->FullyRefine(m_kPrimitives);

	if (m_kPrimitives.GetSize() == 0)
	{
		m_pkLinearNodes = NULL;
		return;
	}

	// Initialize _buildData_ array for primitives
	UT_TVector<BVHPrimitiveInfo> buildData;
	buildData.ReserveCapacity(m_kPrimitives.GetSize());

	for (uint32_t i = 0; i < m_kPrimitives.GetSize(); ++i)
	{
		UT_BBOX_AABB bbox = m_kPrimitives.GetAt(i)->GetWorldBound();
		buildData.AddLast(BVHPrimitiveInfo(i, bbox));
	}

	// Recursively build BVH tree for primitives
	MM_MemoryArena buildArena;
	uint32_t totalNodes = 0;
	UT_TVector<OR_ScenePrimitiveBasePtr> orderedPrims;
	orderedPrims.ReserveCapacity(m_kPrimitives.GetSize());

	BVHBuildNode *root = _recursiveBuild(buildArena, buildData, 0,
		m_kPrimitives.GetSize(), &totalNodes,
		orderedPrims);

	m_kPrimitives.Swap(orderedPrims);

	//Info("BVH created with %d nodes for %d primitives (%.2f MB)", totalNodes,
	//	(int)primitives.size(), float(totalNodes * sizeof(LinearBVHNode))/(1024.f*1024.f));

	// Compute representation of depth-first traversal of BVH tree
	m_pkLinearNodes	 = MM_NEW LinearBVHNode[totalNodes];

	uint32_t offset  = 0;

	_recursiveFlatten(root, &offset);

	UT_ASSERT(offset == totalNodes);
}
//-----------------------------------------------------------------
UT_BBOX_AABB OR_ScenePrimitiveAggregateBVH::GetWorldBound() const
{
	if(NULL == m_pkLinearNodes)
		return UT_BBOX_AABB();
	else
		return m_pkLinearNodes[0].m_kWorldBounds;
}
//-----------------------------------------------------------------
bool OR_ScenePrimitiveAggregateBVH::Intersect(const UT_Ray &ray, OR_Intersection *isect) const
{
	if (NULL == m_pkLinearNodes)
		return false;

	bool hit = false;
	UT_Point3Float origin = ray(ray.mint);
	UT_Vec3Float invDir(1.f / ray.d.x, 1.f / ray.d.y, 1.f / ray.d.z);
	uint32_t dirIsNeg[3] = { invDir.x < 0, invDir.y < 0, invDir.z < 0 };

	// Follow ray through BVH nodes to find primitive intersections
	uint32_t todoOffset = 0, nodeNum = 0;
	uint32_t todo[64];
	while (true) 
	{
		const LinearBVHNode *node = &m_pkLinearNodes[nodeNum];

		// Check ray against BVH node
		if (node->m_kWorldBounds.IntersectP(ray, invDir, dirIsNeg))
		{
			if (node->m_ucPrimitives > 0) 
			{
				// Intersect ray with primitives in leaf BVH node
				for (uint32_t i = 0; i < node->m_ucPrimitives; ++i)
				{
					if (m_kPrimitives.GetAt(node->m_uiPrimitivesOffset+i)->Intersect(ray, isect))
					{
						hit = true;
					}
					else
					{
					
					}
				}
				if (todoOffset == 0) 
					break;
			
				nodeNum = todo[--todoOffset];
			}
			else 
			{
				// Put far BVH node on _todo_ stack, advance to near node
				if (dirIsNeg[node->m_ucAxis]) 
				{
					todo[todoOffset++] = nodeNum + 1;
					nodeNum = node->m_uiSecondChildOffset;
				}
				else 
				{
					todo[todoOffset++] = node->m_uiSecondChildOffset;
					nodeNum = nodeNum + 1;
				}
			}
		}
		else 
		{
			if (todoOffset == 0)
				break;
		
			nodeNum = todo[--todoOffset];
		}
	}

	return hit;
}
//-----------------------------------------------------------------
bool OR_ScenePrimitiveAggregateBVH::IntersectP(const UT_Ray &ray) const
{
	if (NULL == m_pkLinearNodes)
		return false;

	bool hit = false;
	UT_Point3Float origin = ray(ray.mint);
	UT_Vec3Float invDir(1.f / ray.d.x, 1.f / ray.d.y, 1.f / ray.d.z);
	uint32_t dirIsNeg[3] = { invDir.x < 0, invDir.y < 0, invDir.z < 0 };

	// Follow ray through BVH nodes to find primitive intersections
	uint32_t todoOffset = 0, nodeNum = 0;
	uint32_t todo[64];
	while (true) 
	{
		const LinearBVHNode *node = &m_pkLinearNodes[nodeNum];
		if (node->m_kWorldBounds.IntersectP(ray, invDir, dirIsNeg))
		{
			// Process BVH node _node_ for traversal
			if (node->m_ucPrimitives > 0)
			{
				for (uint32_t i = 0; i < node->m_ucPrimitives; ++i)
				{
					if (m_kPrimitives.GetAt(node->m_uiPrimitivesOffset+i)->IntersectP(ray))
					{
						return true;
					}
					else 
					{
					}
				}
				if (todoOffset == 0) 
					break;
				
				nodeNum = todo[--todoOffset];
			}
			else
			{
				if (dirIsNeg[node->m_ucAxis]) 
				{
					/// second child first
					todo[todoOffset++] = nodeNum + 1;
					nodeNum = node->m_uiSecondChildOffset;
				}
				else {
					todo[todoOffset++] = node->m_uiSecondChildOffset;
					nodeNum = nodeNum + 1;
				}
			}
		}
		else
		{
			if (todoOffset == 0)
				break;
	
			nodeNum = todo[--todoOffset];
		}
	}
	return false;
}
//-----------------------------------------------------------------
OR_ScenePrimitiveAggregateBVH::BVHBuildNode* OR_ScenePrimitiveAggregateBVH::_recursiveBuild(MM_MemoryArena& buildArena, 
	UT_TVector<BVHPrimitiveInfo>& buildData, 
	uint32_t start, uint32_t end, PARAM_OUT uint32_t* totalNodes, 
	UT_TVector<OR_ScenePrimitiveBasePtr>& orderedPrims)
{
	UT_ASSERT(start != end);
	(*totalNodes)++;

	BVHBuildNode *node = buildArena.Alloc<BVHBuildNode>();
	
	// Compute bounds of all primitives in BVH node
	UT_BBOX_AABB bbox;
	for (uint32_t i = start; i < end; ++i)
	{
		bbox = UT_MathHelper::Union(bbox, buildData.GetAt(i).m_kWorldBounds);
	}
	
	uint32_t nPrimitives = end - start;
	if (nPrimitives == 1)
	{
		// Create leaf _BVHBuildNode_
		uint32_t firstPrimOffset = orderedPrims.GetSize();
		for (uint32_t i = start; i < end; ++i) 
		{
			uint32_t primNum = buildData.GetAt(i).m_iPrimitiveNumber;
			orderedPrims.AddLast(m_kPrimitives.GetAt(primNum));
		}
		node->InitAsLeaf(firstPrimOffset, nPrimitives, bbox);
	}
	else 
	{
		// Compute bound of primitive centroids, choose split dimension _dim_
		UT_BBOX_AABB centroidBounds;
		for (uint32_t i = start; i < end; ++i)
			centroidBounds = UT_MathHelper::Union(centroidBounds, buildData.GetAt(i).m_kCentroid);

		int dim = centroidBounds.MaximumExtent();

		if (centroidBounds.pMax[dim] == centroidBounds.pMin[dim])
		{
			// Create leaf _BVHBuildNode_
			uint32_t firstPrimOffset = orderedPrims.GetSize();
			for (uint32_t i = start; i < end; ++i) 
			{
				uint32_t primNum = buildData.GetAt(i).m_iPrimitiveNumber;
				orderedPrims.AddLast(m_kPrimitives.GetAt(primNum));
			}

			node->InitAsLeaf(firstPrimOffset, nPrimitives, bbox);
			return node;
		}

		// Partition primitives into two sets and build children
		uint32_t mid = (start + end) / 2;

		node->InitAsInterior(dim,
			_recursiveBuild(buildArena, buildData, start, mid,
			totalNodes, orderedPrims),
			_recursiveBuild(buildArena, buildData, mid, end,
			totalNodes, orderedPrims));
	}
	return node;
}
//-----------------------------------------------------------------
uint32_t OR_ScenePrimitiveAggregateBVH::_recursiveFlatten(BVHBuildNode *node, uint32_t *offset)
{
	LinearBVHNode *linearNode	= &m_pkLinearNodes[*offset];
	linearNode->m_kWorldBounds	= node->m_kWorldBounds;
	uint32_t myOffset = (*offset)++;

	if (node->m_uiPrimitiveNums > 0)
	{
		//Left Node
		UT_ASSERT(!node->m_akChildrenNode[0] && !node->m_akChildrenNode[1]);
		
		linearNode->m_uiPrimitivesOffset = node->m_uiFirstPrimOffset;
		linearNode->m_ucPrimitives		 = node->m_uiPrimitiveNums;
	}
	else
	{
		// Creater interior flattened BVH node
		linearNode->m_ucAxis				= node->m_uiSplitAxis;
		linearNode->m_ucPrimitives			= 0;
		_recursiveFlatten(node->m_akChildrenNode[0], offset);
		linearNode->m_uiSecondChildOffset	= _recursiveFlatten(node->m_akChildrenNode[1],offset);
	}

	return myOffset;
}
//-----------------------------------------------------------------
//-----------------------------------------------------------------
//-----------------------------------------------------------------
//-----------------------------------------------------------------
